import numpy as np
import os
from glob import glob
import cv2
import json
import xml.etree.ElementTree as ET


def get_points(label):
    tree = ET.parse(label)
    root = tree.getroot()
    points = [[int(obj.find('point').find('x').text), int(obj.find('point').find('y').text)]for obj in root.findall('object')]
    return points
    

def generate_data(label_path):
    rgb_path = label_path.replace('/GT_/', '/RGB/').replace('R.xml', '.jpg')
    t_path = label_path.replace('/GT_/', '/Infrared/').replace('R.xml', 'R.jpg')
    rgb = cv2.imread(rgb_path)[..., ::-1].copy()
    t = cv2.imread(t_path)[..., ::-1].copy()
    im_h, im_w, _ = rgb.shape
    print('rgb and t shape', rgb.shape, t.shape)

    points = get_points(label_path)
    points = np.asarray(points)
    # print('points', points.shape)
    idx_mask = (points[:, 0] >= 0) * (points[:, 0] <= im_w) * (points[:, 1] >= 0) * (points[:, 1] <= im_h)
    points = points[idx_mask]
    return rgb, t, points


if __name__ == '__main__':

    root_path = '../../datasets/DroneRGBT/'  # dataset root path
    save_dir = './data/DroneRGBT/'
    # save_dir = '/data2/cjq/RGBT-test'

    for phase in ['Test', 'Train', 'Val']:
        if phase == 'Test':
            save_phase = 'test'
        elif phase == 'Train':
            save_phase = 'train'
        elif phase == 'Val':
            save_phase = 'val'
        else:
            print('ERROR')
            exit()
            
        sub_dir = os.path.join(root_path, phase)
        sub_save_dir = os.path.join(save_dir, save_phase)
        if not os.path.exists(sub_save_dir):
            os.makedirs(sub_save_dir)
        gt_list = glob(os.path.join(sub_dir, 'GT_/*.xml'))
        for gt_path in gt_list:
            name = os.path.basename(gt_path)
            # print('name', name)
            rgb, t, points = generate_data(gt_path)
            
            points = get_points(gt_path)

            pt_json = {"points": points, "count": len(points)}
            

            im_save_path = os.path.join(sub_save_dir, name)
            rgb_save_path = im_save_path.replace('R.xml', '_RGB.jpg')
            t_save_path = im_save_path.replace('R.xml', '_T.jpg')
            cv2.imwrite(rgb_save_path, rgb)
            cv2.imwrite(t_save_path, t)
            
            
            gd_save_path = im_save_path.replace('R.xml', '_GT.json')
            with open(gd_save_path, 'w') as f:
                f.write(json.dumps(pt_json))
